{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# module 导出\n",
    "\n",
    "参考：`tvm/tests/python/runtime/test_runtime_module_export.py`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "import set_env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "temp_dir = Path(\".temp\")\n",
    "temp_dir.mkdir(exist_ok=True, parents=True)\n",
    "# 创建临时目录\n",
    "header_file_dir_path = temp_dir"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `CSourceModuleCreate`\n",
    "\n",
    "```cpp\n",
    "/*!\n",
    " * \\brief Create a C source module for viewing and compiling GCC code.\n",
    " * \\param code The code to be viewed.\n",
    " * \\param fmt The code format.\n",
    " * \\param func_names The name of functions inside the runtime module.\n",
    " * \\param const_vars. The constant variables that the c source module needs.\n",
    " * \\return The created module.\n",
    " */\n",
    "runtime::Module CSourceModuleCreate(const String& code, const String& fmt,\n",
    "                                    const Array<String>& func_names,\n",
    "                                    const Array<String>& const_vars = {});\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "::::{dropdown} CSourceModuleCreate 用于创建 C 源代码模块，以便查看和编译 GCC 代码。\n",
    "\n",
    "参数说明:\n",
    "1. **`const String& code`**  \n",
    "   - 这是字符串类型的参数，表示需要被查看和编译的 GCC 代码。\n",
    "   - `const` 表示这个参数在函数内部不会被修改。\n",
    "   - `String&` 表示这是字符串的引用，避免了不必要的拷贝。\n",
    "\n",
    "2. **`const String& fmt`**  \n",
    "   - 这是字符串类型的参数，表示代码的格式。\n",
    "   - 可能用于指定代码的格式化方式（如缩进、换行等），或者用于指定代码的语言类型（如 C、C++ 等）。\n",
    "\n",
    "3. **`const Array<String>& func_names`**  \n",
    "   - 这是字符串数组类型的参数，表示运行时模块中函数的名称。\n",
    "   - `Array<String>` 可能是自定义的数组类型，用于存储多个字符串。\n",
    "   - 这些函数名称可能是需要在生成的 C 源代码模块中导出的函数。\n",
    "\n",
    "4. **`const Array<String>& const_vars = {}`**  \n",
    "   - 这是可选的字符串数组类型的参数，表示C源代码模块中需要的常量变量。\n",
    "   - `= {}` 表示这个参数有默认值，即空数组。如果调用函数时不提供这个参数，函数会使用空数组作为默认值。\n",
    "\n",
    "返回值:\n",
    "- **`runtime::Module`**  \n",
    "  - 这是函数的返回值类型，表示创建的 C 源代码模块。\n",
    "  - `runtime::Module` 可能是自定义的类或结构体，用于表示运行时模块。\n",
    "\n",
    "函数功能:\n",
    "- 该函数的主要功能是创建 C 源代码模块，该模块可以用于查看和编译 GCC 代码。\n",
    "- 生成的模块可能包含指定的函数名称和常量变量，并且可以根据提供的代码格式进行格式化。\n",
    "\n",
    "使用场景\n",
    "- 这个函数可能用于编译器或代码生成工具中，用于将高级语言代码（如 GCC 代码）转换为 C 源代码模块，以便进一步编译或执行。\n",
    "- 通过指定函数名称和常量变量，可以定制生成的C源代码模块的内容。\n",
    "::::"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "示例调用"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tvm\n",
    "code = r\"int main() { return 0; }\";\n",
    "fmt = \"c\"\n",
    "func_names = [\"main\"]\n",
    "const_vars = [\"MAX_VALUE\"]\n",
    "csource_module = tvm.runtime._ffi_api.CSourceModuleCreate(code, fmt, func_names, const_vars)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "下面介绍更加复杂的例子。\n",
    "\n",
    "先定义一些头文件。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "import textwrap\n",
    "def gen_engine_header(header_file_dir_path):\n",
    "    code = r\"\"\"\n",
    "        #ifndef _ENGINE_H_\n",
    "        #define _ENGINE_H_\n",
    "        #include <cstdint>\n",
    "        #include <string>\n",
    "        #include <sstream>\n",
    "        #include <vector>\n",
    "        class Engine {\n",
    "        };\n",
    "\n",
    "        #endif\n",
    "        \"\"\"\n",
    "    header_file_dir_path = Path(header_file_dir_path)\n",
    "    header_file = header_file_dir_path/\"gcc_engine.h\"\n",
    "    code = textwrap.dedent(code).lstrip()\n",
    "    with open(header_file, \"w\") as f:\n",
    "        f.write(code)\n",
    "\n",
    "def generate_engine_module(header_file_dir_path):\n",
    "    code = r\"\"\"\n",
    "        #include <tvm/runtime/c_runtime_api.h>\n",
    "        #include <dlpack/dlpack.h>\n",
    "        #include \"gcc_engine.h\"\n",
    "\n",
    "        extern \"C\" void gcc_1_(float* gcc_input4, float* gcc_input5,\n",
    "                float* gcc_input6, float* gcc_input7, float* out) {\n",
    "            Engine engine;\n",
    "        }\n",
    "        \"\"\"\n",
    "    import tvm.runtime._ffi_api\n",
    "\n",
    "    gen_engine_header(header_file_dir_path)\n",
    "    code = textwrap.dedent(code).lstrip()\n",
    "    csource_module = tvm.runtime._ffi_api.CSourceModuleCreate(code, \"cc\", [], None)\n",
    "    return csource_module"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 动态库导出"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "One or more operators have not been tuned. Please tune your model for better performance. Use DEBUG logging level to see more details.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "devc.o\tlib0.cc  lib1.c  lib2.o\n"
     ]
    }
   ],
   "source": [
    "from tvm import relay, te\n",
    "import tvm.relay.testing\n",
    "import os\n",
    "os.environ['PATH'] += ':/usr/local/cuda/bin' # 保证 nvcc 可以被找到\n",
    "synthetic_mod, synthetic_params = relay.testing.synthetic.get_workload()\n",
    "synthetic_llvm_mod, synthetic_llvm_params = relay.testing.synthetic.get_workload()\n",
    "with tvm.transform.PassContext(opt_level=3):\n",
    "    synthetic_cpu_lib = relay.build_module.build(\n",
    "        synthetic_llvm_mod, \"llvm\", params=synthetic_llvm_params, mod_name=\"llvmlib\"\n",
    "    )\n",
    "A = te.placeholder((1024,), name=\"A\")\n",
    "B = te.compute(A.shape, lambda *i: A(*i) + 1.0, name=\"B\")\n",
    "s = te.create_schedule(B.op)\n",
    "f = tvm.build(s, [A, B], \"c\", name=\"myadd\")\n",
    "path_lib = temp_dir/\"deploy_lib.so\"\n",
    "engine_module = generate_engine_module(header_file_dir_path)\n",
    "synthetic_cpu_lib.module.import_module(f)\n",
    "synthetic_cpu_lib.module.import_module(engine_module)\n",
    "work_dir = temp_dir/\"work_dir\"\n",
    "work_dir.mkdir(exist_ok=True)\n",
    "kwargs = {\"options\": [\"-O2\", \"-std=c++17\", f\"-I{header_file_dir_path}\"]}\n",
    "synthetic_cpu_lib.export_library(path_lib, fcompile=False, workspace_dir=work_dir, **kwargs)\n",
    "loaded_lib = tvm.runtime.load_module(path_lib)\n",
    "!ls {work_dir} # 查看生成的代码"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 静态库导出"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "生成两个 LLVM 模块："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "A = te.placeholder((1024,), name=\"A\")\n",
    "B = te.compute(A.shape, lambda *i: A(*i) + 1.0, name=\"B\")\n",
    "s = te.create_schedule(B.op)\n",
    "mod0 = tvm.build(s, [A, B], \"llvm\", name=\"myadd0\")\n",
    "mod1 = tvm.build(s, [A, B], \"llvm\", name=\"myadd1\")\n",
    "\n",
    "assert mod0.implements_function(\"myadd0\")\n",
    "assert mod1.implements_function(\"myadd1\")\n",
    "assert mod1.is_dso_exportable"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "``mod1`` 当前是 `'llvm'` 模块，将其保存并重新加载为普通的 `'static_library'`："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "mod1_o_path = f\"{temp_dir}/mod1.o\"\n",
    "mod1.save(mod1_o_path)\n",
    "mod1_o = tvm.runtime.load_static_library(mod1_o_path, [\"myadd1\"])\n",
    "assert mod1_o.implements_function(\"myadd1\")\n",
    "assert mod1_o.is_dso_exportable"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "将 `mod1` 作为静态库导入到 `mod0` 中，并将其编译为独立的 DSO（动态共享对象）。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "mod0.import_module(mod1_o)\n",
    "mod0_dso_path = f\"{temp_dir}/mod0.so\"\n",
    "mod0.export_library(mod0_dso_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "导入的 ``mod1`` 被静态链接到 `mod0` 中。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "loaded_lib = tvm.runtime.load_module(mod0_dso_path)\n",
    "assert loaded_lib.type_key == \"library\"\n",
    "assert len(loaded_lib.imported_modules) == 0\n",
    "assert loaded_lib.implements_function(\"myadd0\")\n",
    "assert loaded_lib.get_function(\"myadd0\")\n",
    "assert loaded_lib.implements_function(\"myadd1\")\n",
    "assert loaded_lib.get_function(\"myadd1\")\n",
    "assert not loaded_lib.is_dso_exportable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
